{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "af6397b2",
   "metadata": {},
   "source": [
    "# Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "839c4a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My OpenAI Key\n",
    "import os\n",
    "import openai\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-....\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d572aa9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hide logs\n",
    "import logging\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.CRITICAL)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "40cf0773",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "### Generate some example Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fa34cd83",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import download_loader\n",
    "from llama_index.indices.vector_store import VectorStoreIndex\n",
    "from llama_index.indices.tree.base import TreeIndex\n",
    "\n",
    "WikipediaReader = download_loader(\"WikipediaReader\")\n",
    "\n",
    "loader = WikipediaReader()\n",
    "documents = loader.load_data(pages=[\"Berlin\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0c32392b",
   "metadata": {},
   "source": [
    "### Create a list of any sort of indices (custom LLMs, custom embeddings, etc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f59e6c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = [\n",
    "    VectorStoreIndex.from_documents(documents),\n",
    "    TreeIndex.from_documents(documents),\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "827ada33",
   "metadata": {},
   "source": [
    "## Using the Playground\n",
    "\n",
    "\n",
    "### Initialize with indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a04e4535",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.playground import Playground\n",
    "\n",
    "playground = Playground(indices=indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5f6999fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mQuery:\u001b[0m\n",
      "What is the population of Berlin?\n",
      "\n",
      "\u001b[1mVectorStoreIndex\u001b[0m, retriever mode = default\n",
      "\u001b[36;1m\u001b[1;3m\n",
      "The population of Berlin is approximately 3.7 million inhabitants.\u001b[0m\n",
      "\n",
      "\u001b[1mTreeIndex\u001b[0m, retriever mode = select_leaf\n",
      "\u001b[33;1m\u001b[1;3m\n",
      "It is not possible to answer this question with the given context information.\u001b[0m\n",
      "\n",
      "\u001b[1mTreeIndex\u001b[0m, retriever mode = select_leaf_embedding\n",
      "\u001b[33;1m\u001b[1;3m\n",
      "The population of Berlin is approximately 3.7 million inhabitants.\u001b[0m\n",
      "\n",
      "\u001b[1mTreeIndex\u001b[0m, retriever mode = all_leaf\n",
      "\u001b[33;1m\u001b[1;3m\n",
      "\n",
      "The population of Berlin is approximately 3.75 million inhabitants. This population has been shaped by the city's turbulent history, with Jewish emigration during the 1930s, the destruction of the city during World War II, and the division of the city into East and West Berlin during the Cold War. Since the reunification of Germany in 1990, Berlin has seen a surge in population growth, with many people from other parts of Germany and the world moving to the city. At the end of 2019, the population of Berlin was estimated to be around 3.75 million inhabitants. The city is home to a diverse religious population, with the faithful of the different religions and denominations maintaining many places of worship in Berlin, including eight parishes of the Independent Evangelical Lutheran Church, 36 Baptist congregations, 29 New Apostolic Churches, 15 United Methodist churches, eight Free Evangelical Congregations, four Churches of Christ, Scientist (1st, 2nd, 3rd, and 11th), six congregations of the Church of Jesus Christ of Latter-day Saints, an Old Catholic church, an Anglican church, more than 80 mosques, ten synagogues, and two Buddhist temples. Berlin is also home to a large number of immigrants from around the world, with 48 percent of the residents under the age of 15 having a migration background in 2017. Berlin is a major economic center in Europe, with many international companies and organizations based in the city, such as the Fraunhofer Society, the Leibniz Association, the Helmholtz Association, and the Max Planck Society, as well as a large number of tourists visiting each year. The city is well-connected to the rest of Germany and Europe through its extensive road, rail, and air transport networks, making it an attractive destination for business and leisure travelers alike. It is also home to a number of renowned research institutions, universities, and medical schools, as well as seven symphony orchestras, including the world-renowned Berlin Philharmonic Orchestra, the Konzerthausorchester Berlin, and the Haus der Kulturen der Welt. Berlin is home to a vibrant cultural and entertainment scene, with a diverse range of cuisine, including Michelin-starred restaurants, vegetarian and vegan offerings, street food, and international cuisine, as well as a variety of botanical gardens, zoos, and other recreational activities. This makes it an attractive destination for people from all over the world. Berlin is also home to two zoos, the Botanischer Garten, the Tiergarten park, and the Gärten der Welt, as well as many cafés, street musicians, beach bars, flea markets, and boutique shops. Berlin has established a high-profile as a host city of major international sporting events, such as the 1936 Summer Olympics, the 2006 FIFA World Cup final, the IAAF World Championships in Athletics, the Basketball Euroleague Final Four, the UEFA Champions League Final, and the 2023 Special Olympics World Summer Games. It is also home to several professional sports teams, such as Hertha BSC, and has a large Olympic training center.\u001b[0m\n",
      "\n",
      "\u001b[1mTreeIndex\u001b[0m, retriever mode = root\n",
      "\u001b[33;1m\u001b[1;3m\n",
      "The population of Berlin is 3.7 million within city limits and 4.5 million in its urban area.\u001b[0m\n",
      "\n",
      "\n",
      "Ran 5 combinations in total.\n"
     ]
    }
   ],
   "source": [
    "result_df = playground.compare(\"What is the population of Berlin?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0ad183e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Index</th>\n",
       "      <th>Retriever Mode</th>\n",
       "      <th>Output</th>\n",
       "      <th>Duration</th>\n",
       "      <th>Prompt Tokens</th>\n",
       "      <th>Completion Tokens</th>\n",
       "      <th>Embed Tokens</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>VectorStoreIndex</td>\n",
       "      <td>default</td>\n",
       "      <td>\\nThe population of Berlin is approximately 3....</td>\n",
       "      <td>2.525580</td>\n",
       "      <td>1786</td>\n",
       "      <td>13</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>select_leaf</td>\n",
       "      <td>\\nIt is not possible to answer this question w...</td>\n",
       "      <td>5.536037</td>\n",
       "      <td>4732</td>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>select_leaf_embedding</td>\n",
       "      <td>\\nThe population of Berlin is approximately 3....</td>\n",
       "      <td>5.426232</td>\n",
       "      <td>897</td>\n",
       "      <td>13</td>\n",
       "      <td>9146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>all_leaf</td>\n",
       "      <td>\\n\\nThe population of Berlin is approximately ...</td>\n",
       "      <td>238.278128</td>\n",
       "      <td>27291</td>\n",
       "      <td>5035</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>root</td>\n",
       "      <td>\\nThe population of Berlin is 3.7 million with...</td>\n",
       "      <td>3.375349</td>\n",
       "      <td>558</td>\n",
       "      <td>23</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              Index         Retriever Mode  \\\n",
       "0  VectorStoreIndex                default   \n",
       "1         TreeIndex            select_leaf   \n",
       "2         TreeIndex  select_leaf_embedding   \n",
       "3         TreeIndex               all_leaf   \n",
       "4         TreeIndex                   root   \n",
       "\n",
       "                                              Output    Duration  \\\n",
       "0  \\nThe population of Berlin is approximately 3....    2.525580   \n",
       "1  \\nIt is not possible to answer this question w...    5.536037   \n",
       "2  \\nThe population of Berlin is approximately 3....    5.426232   \n",
       "3  \\n\\nThe population of Berlin is approximately ...  238.278128   \n",
       "4  \\nThe population of Berlin is 3.7 million with...    3.375349   \n",
       "\n",
       "   Prompt Tokens  Completion Tokens  Embed Tokens  \n",
       "0           1786                 13             7  \n",
       "1           4732                115             0  \n",
       "2            897                 13          9146  \n",
       "3          27291               5035             0  \n",
       "4            558                 23             0  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result_df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8829a829",
   "metadata": {},
   "source": [
    "### Initialize with Documents\n",
    "\n",
    "Automatically construct the playground using a vector, tree, and summary index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbc8ade",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uses documents in a preset list of indices\n",
    "playground = Playground.from_docs(documents=documents)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
